#!/user/bin/env python3
# -*- coding: utf-8 -*-
import time
import calendar as c
from datetime import datetime, date, timedelta
from dateutil.relativedelta import relativedelta
# from workalendar.asia import china as cn
from numpy import long
# from .MySQLDBUtil import read_data
import pandas as pd
from jqdata import *


# count:提前多少天，int
# date_str '%Y%m%d' same as tushare data
def getPreviousDay(date_str: str, count: int) -> str:
    # today = date.today()
    date_time = datetime.strptime(date_str, '%Y-%m-%d')
    daycount = timedelta(days=count)
    prvdate = date_time - daycount
    previous_str = prvdate.strftime('%Y-%m-%d')
    return previous_str


# 判断某天是否交易日，不是的话迁移到下一个交易日
def switchToNextTradeDay(somedate: str)->str:
    tradedate_l = get_trade_days(start_date=somedate)

    if somedate not in tradedate_l:
        return getPreviousDay(somedate, -1)
    else:
        return somedate


def createStartEndDates(year):
    return str(year) + '0101', str(year) + '1231'


# months: 相对当前时间的月数
def getFutureMonth(months: int, from_date=None)->datetime:
    if from_date is None:
        from_date = datetime.now()
    fudate = from_date + relativedelta(months=months)
    return fudate


def toStringYearMonth(somedate, n: int) -> str :
    fudate = getFutureMonth(n, somedate)
    return datetime.strftime(fudate, '%Y-%m-%d')


# 计算相差月份
def calcMonthsLeg(startdate: date, enddate: date):
    samemonthdate = None
    try:
        samemonthdate = date(enddate.year, enddate.month, startdate.day)
    except Exception as e:
        print('计算相差月份出错')
        samemonthdate = date(enddate.year, enddate.month, c.monthrange(enddate.year, enddate.month[1]))
        
    holdmonths = 0
    decimalmonth = 0.0

    if samemonthdate > enddate:
        premonthdate = None
        try:
            premonthdate = date(enddate.year, enddate.month - 1, startdate.day)
        except Exception as e:
            print(e)
            premonthdate = date(enddate.year, enddate.month - 1, c.monthrange(enddate.year, enddate.month - 1)[1])
        currmonthdays = (samemonthdate - premonthdate).days
        holdmonths = (premonthdate.year - startdate.year
                      ) * 12 + premonthdate.month - startdate.month
        decimalmonth = (enddate - premonthdate).days / currmonthdays
    elif samemonthdate < enddate:
        nextmonthdate = None
        try:
            nextmonthdate = date(enddate.year, enddate.month + 1, startdate.day)
        except Exception as e:
            nextmonthdate = date(enddate.year, enddate.month + 1, c.monthrange(enddate.year, enddate.month + 1)[1])
        currmonthdays = (nextmonthdate - samemonthdate).days
        holdmonths = (samemonthdate.year - startdate.year
                      ) * 12 + samemonthdate.month - startdate.month
        decimalmonth = (enddate - samemonthdate).days / currmonthdays
    else:
        holdmonths = (enddate.year - startdate.year) * 12 + enddate.month - startdate.month

    return holdmonths, decimalmonth


# 返回n年前的日期
def calcYearLeg(years, from_date=None):
    if from_date is None:
        from_date = datetime.now()
    return from_date - relativedelta(years=years)

#
def getFutureMonth(months: int, from_date=None):
    if from_date is None:
        from_date = datetime.now()
    fudate = from_date + relativedelta(months=months)
    return fudate

# 获取两个日期间所有月份
def get_month_range(start_day, end_day):
  months_leg = (end_day.year - start_day.year)*12 + end_day.month - start_day.month
  month_range = ['%s-%02d-%02d'%(start_day.year + mon//12,mon%12+1, 1)
                    for mon in range(start_day.month-1,start_day.month + months_leg)]
  return month_range


# 利用pandas.date_range()函数，稍加改造返回年份的字符串list，本年没结束的话不包含本年
# freq ref to https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#timeseries-offset-aliases
def getYearLeg(start_day: str, end_day: str)->list:
    date = pd.date_range(start_day, end_day, freq='Y')
    year_list = [datetime.strftime(x, '%F') for x in date]

    return year_list


# 计算函数耗时
def time_me(func):
    def _wrapper(*args, **kwargs):
        start = time.clock()
        func(*args, **kwargs)
        print( '%s cost %s second'%(func.__name__, time.clock()))
        return _wrapper()

# 日期格式变换
# sourceformat 源日期格式 target 目标日期格式
def dateformatchange(date_str, sourceformat, targetformat):
    t1 = time.strptime(date_str, sourceformat)
    chg = time.strftime(targetformat, t1)
    # print('date',chg)
    return (chg)

# def latestWorkingDay(aday):
    # return cn.China.find_following_working_day(day = aday)


# def getWorkingDaysInPeriod(startdate,enddate):
#     date_list = []
#     # cal = cn.China()
#     cal = ql.China(ql.China.IB)
#     # begin_date = datetime.datetime.strptime(startdate, "%Y-%m-%d")
#     # end_date = datetime.datetime.strptime(time.strftime('%Y-%m-%d', time.localtime(time.time())), "%Y-%m-%d")
#     while startdate <= enddate:
#         if cal.is_working_day(startdate):
#             date_str = startdate.strftime("%Y%m%d")
#             date_list.append(date_str)
#         startdate += timedelta(days=1)
#     return date_list


# 输入日期格式-YYYYMMDD string,若为节假日，往前直至上一工作日。
# def handleHoliday(exchange, date_str):
#     sql = 'SELECT is_open FROM tushare.trade_cal WHERE cal_date = "' + date_str +'" and exchange = "' + exchange + '"'
#     dtdf = read_data(sql)
#     dtlist = dtdf.values.tolist()
#
#     #在交易日历内
#     if len(dtlist):
#         flag = ''.join(dtlist[0])
#         if flag == '1': #当日为工作日
#             return date_str
#         else:
#             prv = getPreviousDay(date_str,1)
#             handleHoliday(exchange,prv)
#     else:
#         print(date_str,'不在交易日历内')


#获取N天前/后的交易日日期
# shift < 0往前取， >0 往后取
def shiftTradingDay(adate: datetime.date, shift=5)-> datetime.date:
    # 获取所有的交易日，返回一个包含所有交易日的 list,元素值为 datetime.date 类型.

    # get_all_trade_days()
    # 得到date之后shift天那一天在列表中的行标号 返回一个数
    trade_cal = getTradeDaysWithCount(datetime.strftime(adate, '%Y-%m-%d'), shift).tolist()

    return trade_cal[len(trade_cal)-1]


# 入参与上面的函数不同，返回的数据类型也不一样。先不合并
def OffsetDate(end_date: str, count: int) -> datetime.date:
    '''
    end_date:为基准日期
    count:为正则后推，负为前推
    -----------
    return datetime.date
    '''
    trade_date = get_trade_days(end_date=end_date, count=1)[0] # online

    if count > 0:
        # 将end_date转为交易日
        trade_cal = getAllTradeDays().tolist()
            # get_all_trade_days().tolist()
        trade_idx = trade_cal.index(trade_date)

        return trade_cal[trade_idx + count]

    elif count < 0:
        return get_trade_days(end_date=trade_date, count=abs(count))[0]
    else:
        raise ValueError('别闹！')


# 获取年末季末时点
def GetTradePeriod(start_date: str, end_date: str, freq: str = 'ME') -> list:
    '''
    start_date/end_date:str YYYY-MM-DD
    freq:M月，Q季,Y年 默认ME E代表期末 S代表期初
    ================
    return  list[datetime.date]
    '''
    # days = pd.Index(pd.to_datetime(get_trade_days(start_date, end_date)))
    days = pd.Index(pd.to_datetime(getTradeDays(start_date, end_date)))
    idx_df = days.to_frame()

    if freq[-1] == 'E':
        day_range = idx_df.resample(freq[0]).last()
    else:
        day_range = idx_df.resample(freq[0]).first()

    day_range = day_range[0].dt.date

    return day_range.dropna().values.tolist()

# 跟上面yearleg不同的是，本函数包含本年
def getYearsBetweenDates(start_date: str, end_date: str):
    startDate = datetime.strptime(start_date, '%Y-%m-%d').date()
    endDate = datetime.strptime(end_date, '%Y-%m-%d').date()
    return list(range(startDate.year, endDate.year + 1))


def getYearMonthBetweenDates(start_date: str, end_date: str):
    '''获取时间段内年月份'''
    startdate = datetime.strptime(str(start_date), '%Y-%m-%d')
    endtime = datetime.strptime(str(end_date), '%Y-%m-%d')

    delta = startdate.month

    def add_months(dt, months):
        ''' 月份相加 '''
        targetmonth = months + dt.month
        try:
            dt = dt.replace(year=dt.year + int(targetmonth / 12), month=(targetmonth % 12))
        except:
            dt = dt.replace(year=dt.year + int((targetmonth + 1) / 12), month=((targetmonth + 1) % 12), day=1)
            dt += timedelta(days=-1)
        return dt

    n = 0
    date_list = []
    while 1:
        if startdate <= endtime:
            year_month = add_months(startdate, n)
            n = n + 1

            month_str = str(year_month.month)

            if len(month_str) == 1:
                month_str = '0' + month_str

            date_list.append(str(year_month.year) + '-' + month_str)

            if year_month.year == endtime.year and year_month.month == endtime.month:
                break
    return date_list


# beginDate：区间起始日，beginDate默认为当月首日加入list; endDate： 区间结束日
#
def months_s_e_dates(beginDate, endDate):
    date_index = pd.date_range(beginDate, endDate)
    days = [pd.Timestamp(x).strftime("%Y-%m-%d") for x in date_index.values]
    list_month_dates = []
    month_begin_l = month_end_l = []
    for index, v in enumerate(days):
        if index == len(days) - 1:
            list_month_dates.append(days[index])
        if index == 0:
            list_month_dates.append(days[0])
        else:
            _ = v.split('-')[2]
            if _ == '01':
                print(days[index - 1])
                list_month_dates.append(days[index - 1])
                list_month_dates.append(days[index])

    return list_month_dates


if __name__ == '__main__':
    # getYearLeg()
    print(months_s_e_dates('2010-01-02', '2010-04-01'))

